package com.njupt.wuaiagent.agent;

import cn.hutool.core.util.StrUtil;
import com.njupt.wuaiagent.agent.model.AgentState;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CompletableFuture;

/**
 * @Author: wujiaming
 * @CreateTime: 2025/5/22 19:55
 * @Description:抽象类代理类，用类模板方法设计模式，定义了执行流程，用于管理状态的管理，事件的循环AgentLoop
 * 提供了状态的转换、内存管理和基于步骤的执行循环基本功能。子类必须实现step
 * @Version: 1.0
 */


@Data
@Slf4j
public abstract class BaseAgent {

    //1、智能体的名字
    private String name;

    //2、prompt提示词
    private String systemPrompt;
    private String nextStepPrompt;  //AI完成一部操作之后需要下一步操作，为了引导他完成下一步操作要输入给AI的提示词

    //3、状态
    private AgentState state = AgentState.IDLE;

    //4、执行控制
    private int maxStep = 10;
    private int currentStep = 0;

    //5、传入的大模型
    private ChatClient chatClient;

    //6、Memory（需要自己维护的会话上下文）
    private List<Message> messageList = new ArrayList<>();

    /**
     * 运行代理
     *
     * @param userPrompt 用户提示词
     * @return 执行结果
     */
    public String run(String userPrompt){
        //1、幂等判断
        if(this.state != AgentState.IDLE){
            throw new RuntimeException("Cannot run agent from state: " + this.state);
        }
        if(StrUtil.isBlank(userPrompt)){
            throw new RuntimeException("Cannot run agent with empty user prompt");
        }

        //2、修改状态
        this.state = AgentState.RUNNING;

        //3、记录上下文信息
        messageList.add(new UserMessage(userPrompt));


        //4、保存每一个单步执行的执行的结果
        List<String> results = new ArrayList<>();
        try {
            for(int i = 0 ; i < maxStep && AgentState.FINISHED != state; i++){
                int stepNumber = i + 1;
                currentStep = stepNumber;
                log.info("Executing step " + stepNumber + "/" + maxStep);
                //单步执行
                String stepResult = step();
                String result = "Step " + stepNumber + ": " + stepResult;
                results.add(result);
            }
            //检查是否晁补步骤限制
            if(currentStep >= maxStep){
                state = AgentState.FINISHED;
                results.add("Terminated: Reached max steps (" + maxStep + ")");
            }
            return String.join("\n",results);
        } catch (Exception e) {
            state = AgentState.ERROR;
            log.error("Error executing agent",e);
            return "执行错误" + e.getMessage();
        } finally {
            cleanup();
        }

    }


    /**
     * 运行代理
     *
     * @param userPrompt 用户提示词
     * @return 执行结果
     */
    public SseEmitter runStream(String userPrompt){
        SseEmitter sseEmitter = new SseEmitter(300000L);

        CompletableFuture.runAsync(()->{
            //1、幂等判断
            try {
                if(this.state != AgentState.IDLE){
                    sseEmitter.send("错误：无法从运行状态代理 "+ this.state);
                    sseEmitter.complete();
                    return ;
                }
                if(StrUtil.isBlank(userPrompt)){
                    sseEmitter.send("错误：不能使用空提示词进行代理 ");
                    sseEmitter.complete();
                    return ;
                }
            } catch (IOException e) {
                sseEmitter.completeWithError(e);
            }

            //2、修改状态
            this.state = AgentState.RUNNING;

            //3、记录上下文信息
            messageList.add(new UserMessage(userPrompt));


            //4、保存每一个单步执行的执行的结果
            List<String> results = new ArrayList<>();
            try {
                for(int i = 0 ; i < maxStep && AgentState.FINISHED != state; i++){
                    int stepNumber = i + 1;
                    currentStep = stepNumber;
                    log.info("Executing step " + stepNumber + "/" + maxStep);
                    //单步执行
                    String stepResult = step();
                    String result = "Step " + stepNumber + ": " + stepResult;
//                    results.add(result);

                    //发送每一步的结果
                    sseEmitter.send(result);
                }
                //检查是否晁补步骤限制
                if(currentStep >= maxStep){
                    state = AgentState.FINISHED;
                    sseEmitter.send("执行结束达到最大步骤（" + maxStep + ")");
//                    results.add("Terminated: Reached max steps (" + maxStep + ")");
                }
                //正常完成
                sseEmitter.complete();
            } catch (Exception e) {
                state = AgentState.ERROR;
                log.error("执行能执行失败",e);
                try {
                    sseEmitter.send("执行错误" + e.getMessage());
                    sseEmitter.complete();
                } catch (IOException ex) {
                    sseEmitter.completeWithError(e);
                }

            } finally {
                cleanup();
            }
        });

        //设置超时回调
        sseEmitter.onTimeout(()->{
            this.state = AgentState.ERROR;
            this.cleanup();
            log.warn("SSE connection time out");
        });

        //设置完成回调
        sseEmitter.onCompletion(()->{
            if(AgentState.RUNNING == this.state){
                this.state = AgentState.FINISHED;
            }
            this.cleanup();
            log.info("SSE connection completed");
        });

        return sseEmitter;
    }

    /**
     * 执行单个步骤
     * @return
     */
    public abstract String step();


    /**
     * 清理资源
     */
    protected void cleanup(){
        //子类可以重写此方法来清理资源
    }
}
